"""
文档搜索工具模块

提供文档检索、搜索和内容提取功能
"""

import os
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional

from .base_tool import BaseTool, ToolResult


@dataclass
class Document:
    """文档对象"""
    path: str
    title: str
    content: str
    file_type: str
    size: int
    modified_time: float
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


class DocumentSearchTool(BaseTool):
    """
    文档搜索工具
    
    提供文档检索、搜索和内容提取功能
    """
    
    def __init__(self, search_paths: Optional[List[str]] = None):
        """
        初始化文档搜索工具
        
        Args:
            search_paths: 搜索路径列表，如果为None则使用当前目录
        """
        super().__init__(
            name="document_search_tool",
            description="搜索和检索文档内容，支持多种文件格式",
        )
        
        self.search_paths = search_paths or [str(Path.cwd())]
        self.supported_extensions = {
            '.txt', '.md', '.rst', '.py', '.js', '.html', '.xml', 
            '.json', '.yaml', '.yml', '.csv', '.log'
        }
        self._document_cache = {}
        self._cache_ttl = 300  # 5分钟缓存
    
    def execute(self, operation: str, **kwargs) -> ToolResult:
        """
        执行文档搜索操作
        
        Args:
            operation: 操作类型 (search, find, extract, list, info)
            **kwargs: 其他参数
            
        Returns:
            操作结果
        """
        try:
            if operation == "search":
                return self._search_content(**kwargs)
            elif operation == "find":
                return self._find_files(**kwargs)
            elif operation == "extract":
                return self._extract_content(**kwargs)
            elif operation == "list":
                return self._list_documents(**kwargs)
            elif operation == "info":
                return self._get_document_info(**kwargs)
            elif operation == "index":
                return self._index_documents(**kwargs)
            else:
                return ToolResult(
                    success=False,
                    error_message=f"不支持的操作类型: {operation}"
                )
                
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"文档搜索操作失败: {str(e)}"
            )
    
    def _search_content(self, query: str, path: Optional[str] = None, 
                       file_type: Optional[str] = None, limit: int = 10,
                       case_sensitive: bool = False, regex: bool = False,
                       **kwargs) -> ToolResult:
        """搜索文档内容"""
        search_paths = [path] if path else self.search_paths
        results = []
        
        for search_path in search_paths:
            if not os.path.exists(search_path):
                continue
            
            # 遍历目录
            for root, dirs, files in os.walk(search_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    
                    # 检查文件类型
                    if file_type and not file.endswith(file_type):
                        continue
                    
                    # 检查文件扩展名
                    if not self._is_supported_file(file_path):
                        continue
                    
                    try:
                        # 搜索文件内容
                        matches = self._search_in_file(
                            file_path, query, case_sensitive, regex
                        )
                        
                        if matches:
                            results.extend(matches)
                            
                            if len(results) >= limit:
                                break
                    
                    except Exception as e:
                        continue
                
                if len(results) >= limit:
                    break
            
            if len(results) >= limit:
                break
        
        # 限制结果数量
        results = results[:limit]
        
        return ToolResult(
            success=True,
            data=results,
            metadata={
                "query": query,
                "total_matches": len(results),
                "search_paths": search_paths,
                "case_sensitive": case_sensitive,
                "regex": regex
            }
        )
    
    def _find_files(self, pattern: str, path: Optional[str] = None,
                   file_type: Optional[str] = None, limit: int = 50,
                   **kwargs) -> ToolResult:
        """查找文件"""
        search_paths = [path] if path else self.search_paths
        results = []
        
        for search_path in search_paths:
            if not os.path.exists(search_path):
                continue
            
            # 遍历目录
            for root, dirs, files in os.walk(search_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    
                    # 检查文件类型
                    if file_type and not file.endswith(file_type):
                        continue
                    
                    # 检查文件名模式
                    if pattern.lower() in file.lower():
                        try:
                            doc_info = self._get_file_info(file_path)
                            results.append(doc_info)
                            
                            if len(results) >= limit:
                                break
                        except Exception:
                            continue
                
                if len(results) >= limit:
                    break
            
            if len(results) >= limit:
                break
        
        return ToolResult(
            success=True,
            data=results,
            metadata={
                "pattern": pattern,
                "total_files": len(results),
                "search_paths": search_paths
            }
        )
    
    def _extract_content(self, path: str, start_line: Optional[int] = None,
                        end_line: Optional[int] = None, **kwargs) -> ToolResult:
        """提取文档内容"""
        if not os.path.exists(path):
            return ToolResult(
                success=False,
                error_message=f"文件不存在: {path}"
            )
        
        try:
            with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                lines = f.readlines()
            
            # 提取指定行范围
            if start_line is not None or end_line is not None:
                start = start_line or 1
                end = end_line or len(lines)
                start = max(0, start - 1)  # 转换为0基索引
                end = min(len(lines), end)
                content = ''.join(lines[start:end])
            else:
                content = ''.join(lines)
            
            # 获取文件信息
            doc_info = self._get_file_info(path)
            doc_info.content = content
            
            return ToolResult(
                success=True,
                data=doc_info,
                metadata={
                    "operation": "extract",
                    "path": path,
                    "start_line": start_line,
                    "end_line": end_line,
                    "content_length": len(content)
                }
            )
            
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"提取内容失败: {str(e)}"
            )
    
    def _list_documents(self, path: Optional[str] = None, 
                       file_type: Optional[str] = None, limit: int = 100,
                       **kwargs) -> ToolResult:
        """列出文档"""
        search_paths = [path] if path else self.search_paths
        results = []
        
        for search_path in search_paths:
            if not os.path.exists(search_path):
                continue
            
            # 遍历目录
            for root, dirs, files in os.walk(search_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    
                    # 检查文件类型
                    if file_type and not file.endswith(file_type):
                        continue
                    
                    # 检查文件扩展名
                    if not self._is_supported_file(file_path):
                        continue
                    
                    try:
                        doc_info = self._get_file_info(file_path)
                        results.append(doc_info)
                        
                        if len(results) >= limit:
                            break
                    except Exception:
                        continue
                
                if len(results) >= limit:
                    break
            
            if len(results) >= limit:
                break
        
        return ToolResult(
            success=True,
            data=results,
            metadata={
                "operation": "list",
                "total_documents": len(results),
                "search_paths": search_paths
            }
        )
    
    def _get_document_info(self, path: str, **kwargs) -> ToolResult:
        """获取文档信息"""
        if not os.path.exists(path):
            return ToolResult(
                success=False,
                error_message=f"文件不存在: {path}"
            )
        
        try:
            doc_info = self._get_file_info(path)
            
            return ToolResult(
                success=True,
                data=doc_info,
                metadata={
                    "operation": "info",
                    "path": path
                }
            )
            
        except Exception as e:
            return ToolResult(
                success=False,
                error_message=f"获取文档信息失败: {str(e)}"
            )
    
    def _index_documents(self, path: Optional[str] = None, 
                        force_refresh: bool = False, **kwargs) -> ToolResult:
        """索引文档"""
        search_paths = [path] if path else self.search_paths
        indexed_count = 0
        
        for search_path in search_paths:
            if not os.path.exists(search_path):
                continue
            
            # 遍历目录
            for root, dirs, files in os.walk(search_path):
                for file in files:
                    file_path = os.path.join(root, file)
                    
                    # 检查文件扩展名
                    if not self._is_supported_file(file_path):
                        continue
                    
                    # 检查缓存
                    if not force_refresh and file_path in self._document_cache:
                        cache_time = self._document_cache[file_path].get('timestamp', 0)
                        if time.time() - cache_time < self._cache_ttl:
                            continue
                    
                    try:
                        doc_info = self._get_file_info(file_path)
                        self._document_cache[file_path] = {
                            'document': doc_info,
                            'timestamp': time.time()
                        }
                        indexed_count += 1
                    except Exception:
                        continue
        
        return ToolResult(
            success=True,
            data={
                "indexed_count": indexed_count,
                "cache_size": len(self._document_cache)
            },
            metadata={
                "operation": "index",
                "search_paths": search_paths,
                "force_refresh": force_refresh
            }
        )
    
    def _search_in_file(self, file_path: str, query: str, 
                       case_sensitive: bool, regex: bool) -> List[Dict]:
        """在文件中搜索内容"""
        matches = []
        
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                lines = f.readlines()
            
            # 准备搜索模式
            if regex:
                try:
                    pattern = re.compile(query if case_sensitive else query.lower())
                except re.error:
                    return matches
            else:
                search_text = query if case_sensitive else query.lower()
            
            # 搜索每一行
            for line_num, line in enumerate(lines, 1):
                line_text = line if case_sensitive else line.lower()
                
                if regex:
                    if pattern.search(line_text):
                        matches.append({
                            "file_path": file_path,
                            "line_number": line_num,
                            "line_content": line.strip(),
                            "match_type": "regex"
                        })
                else:
                    if search_text in line_text:
                        # 高亮匹配内容
                        highlighted = self._highlight_match(line, query, case_sensitive)
                        matches.append({
                            "file_path": file_path,
                            "line_number": line_num,
                            "line_content": line.strip(),
                            "highlighted_content": highlighted,
                            "match_type": "text"
                        })
        
        except Exception:
            pass
        
        return matches
    
    def _highlight_match(self, line: str, query: str, case_sensitive: bool) -> str:
        """高亮匹配内容"""
        if case_sensitive:
            return line.replace(query, f"**{query}**")
        else:
            # 不区分大小写的高亮
            import re
            pattern = re.compile(re.escape(query), re.IGNORECASE)
            return pattern.sub(lambda m: f"**{m.group()}**", line)
    
    def _get_file_info(self, file_path: str) -> Document:
        """获取文件信息"""
        stat = os.stat(file_path)
        file_ext = Path(file_path).suffix.lower()
        
        # 获取文件标题（文件名或第一行内容）
        title = Path(file_path).stem
        
        # 尝试读取第一行作为标题
        try:
            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                first_line = f.readline().strip()
                if first_line and len(first_line) < 100:
                    title = first_line
        except Exception:
            pass
        
        return Document(
            path=file_path,
            title=title,
            content="",  # 内容在需要时再读取
            file_type=file_ext,
            size=stat.st_size,
            modified_time=stat.st_mtime,
            metadata={
                "readable": True,
                "encoding": "utf-8"
            }
        )
    
    def _is_supported_file(self, file_path: str) -> bool:
        """检查是否为支持的文件类型"""
        file_ext = Path(file_path).suffix.lower()
        return file_ext in self.supported_extensions
    
    def clear_cache(self) -> None:
        """清除缓存"""
        self._document_cache.clear()
    
    def get_cache_stats(self) -> Dict[str, Any]:
        """获取缓存统计信息"""
        return {
            "cache_size": len(self._document_cache),
            "cache_ttl": self._cache_ttl,
            "supported_extensions": list(self.supported_extensions)
        }


class DocumentIndexer:
    """文档索引器"""
    
    def __init__(self, search_tool: DocumentSearchTool):
        """
        初始化文档索引器
        
        Args:
            search_tool: 文档搜索工具实例
        """
        self.search_tool = search_tool
        self.index = {}
    
    def build_index(self, paths: Optional[List[str]] = None) -> Dict[str, Any]:
        """
        构建文档索引
        
        Args:
            paths: 要索引的路径列表
            
        Returns:
            索引统计信息
        """
        result = self.search_tool.execute("index", path=paths[0] if paths else None)
        
        if result.success:
            # 构建倒排索引
            self._build_inverted_index()
            
            return {
                "indexed_documents": result.data["indexed_count"],
                "index_terms": len(self.index),
                "cache_size": result.data["cache_size"]
            }
        
        return {"error": "索引构建失败"}
    
    def _build_inverted_index(self) -> None:
        """构建倒排索引"""
        self.index = {}
        
        for file_path, cache_data in self.search_tool._document_cache.items():
            doc = cache_data['document']
            
            # 简单的词频统计
            words = re.findall(r'\b\w+\b', doc.title.lower())
            for word in words:
                if word not in self.index:
                    self.index[word] = []
                self.index[word].append({
                    "file_path": file_path,
                    "title": doc.title,
                    "type": "title"
                })
    
    def search_index(self, query: str) -> List[Dict[str, Any]]:
        """
        在索引中搜索
        
        Args:
            query: 搜索查询
            
        Returns:
            搜索结果
        """
        results = []
        query_words = re.findall(r'\b\w+\b', query.lower())
        
        for word in query_words:
            if word in self.index:
                results.extend(self.index[word])
        
        # 去重并按相关性排序
        seen = set()
        unique_results = []
        for result in results:
            key = result['file_path']
            if key not in seen:
                seen.add(key)
                unique_results.append(result)
        
        return unique_results


# 全局文档搜索工具实例
_doc_search_tool = DocumentSearchTool()


def get_document_search_tool() -> DocumentSearchTool:
    """获取文档搜索工具的便捷函数"""
    return _doc_search_tool


def search_documents(query: str, **kwargs) -> ToolResult:
    """搜索文档的便捷函数"""
    return _doc_search_tool.execute("search", query=query, **kwargs)


def find_files(pattern: str, **kwargs) -> ToolResult:
    """查找文件的便捷函数"""
    return _doc_search_tool.execute("find", pattern=pattern, **kwargs)
